package cn.itcast.openai;

import com.openai.client.OpenAIClient;
import com.openai.client.okhttp.OpenAIOkHttpClient;
import com.openai.models.chat.completions.ChatCompletionAssistantMessageParam;
import com.openai.models.chat.completions.ChatCompletionCreateParams;
import com.openai.models.chat.completions.ChatCompletionMessageParam;
import com.openai.models.chat.completions.ChatCompletionUserMessageParam;

import java.util.ArrayList;
import java.util.List;
import java.util.Optional;

/**
 * 多轮对话Demo
 */
public class CompletionsMultipleRoundsDemo {

    private static OpenAIClient client;

    public static void main(String[] args) {
        // 创建客户端，指定 API Key 与 baseUrl，其中API KEY从系统环境变量中获取
        client = OpenAIOkHttpClient.builder()
                .baseUrl("https://dashscope.aliyuncs.com/compatible-mode/v1")
                //setx ALIYUN_API_KEY "sk-ef2d238d7dce4505b2d8aec49048b8ee"
                .apiKey(System.getenv("ALIYUN_API_KEY"))
                .build();
        // 创建消息集合，用于存储对话历史记录
        var messageParams = new ArrayList<ChatCompletionMessageParam>();
        // 第一次对话
        chat("我叫nbspdd，请记住我", messageParams);
        System.out.println("------------------------");
        // 第二次对话
        chat("我是谁？", messageParams);
    }

    public static void chat(String userMessage, List<ChatCompletionMessageParam> messageList) {
        // 添加用户消息
        messageList.add(ChatCompletionMessageParam.ofUser(ChatCompletionUserMessageParam.builder().content(userMessage).build()));
        // 构造聊天参数
        var createParams = ChatCompletionCreateParams.builder()
                .model("qwen-plus") // 指定模型
                .messages(messageList) // 添加用户消息
                .build();

        // 调用接口，获取结果并打印
        client.chat().completions()
                .create(createParams)
                .choices()
                .stream()
                .flatMap(choice -> {
                    // 获取 assistant 消息
                    Optional<String> contentOptional = choice.message().content();
                    // 如果有 assistant 消息，则手动构建 assistant 消息对象，并且放到消息集合中
                    if (contentOptional.isPresent()) {
                        // 手动构建 assistant 消息对象，并且放到消息集合中
                        ChatCompletionAssistantMessageParam assistantMessageParam = ChatCompletionAssistantMessageParam.builder()
                                .content(contentOptional.get())
                                .build();
                        messageList.add(ChatCompletionMessageParam.ofAssistant(assistantMessageParam));
                    }
                    // 返回 assistant 消息流
                    return contentOptional.stream();
                })
                .forEach(System.out::println);
    }

}
